'''
Created on Mar 6, 2012

@author: oabalbin
'''
'''
Created on Feb 8, 2012

@author: oabalbin
'''
import os
import sys
import glob
import subprocess
from collections import defaultdict

JOB_ERROR=1
JOB_SUCCESS=0

def create_project(rootdir,name,sample_dict,build,vsannfile,varscan=True):
    '''
    vtools select TCGA_CH_5748 'vsvcf.STATUS_INFO="Somatic"' -t somatic
    vtools use ../vsvcf.ann --files ../TCGA-EJ-5514.somatic.snps.vars.raw.2.vcf ../TCGA-CH-5748.somatic.snps.vars.raw.2.vcf
   '''    
    os.chdir(rootdir)
    args=['mkdir',name]
    retcode = subprocess.call(args)
    os.chdir(os.path.join(rootdir,name))
        
    args=['vtools','init',name]
    retcode=subprocess.call(args)

    if retcode!=0:
        return JOB_ERROR
    cmd=['vtools','import']
    for sp,file in sample_dict.iteritems():
        args=cmd+[file]+['--sample_name',sp]+['--build',build]
        if retcode != 0:
            raise JOB_ERROR
        else:
            print args
            retcode=subprocess.call(args)
    
    # load the varscan-vcf annotation.
    if retcode!=0:
        return JOB_ERROR
    elif varscan:
        args=['vtools','use',vsannfile,'--files']+sample_dict.values() # all snv files to create a giant annotation file of all mutations.
        retcode=subprocess.call(args)
        if retcode != 0:
            raise JOB_ERROR        
        load_all_annotation(rootdir,name)
        
    return JOB_SUCCESS


def create_children_project(rootdir,name,sample_dict,build,vsannfile,annotation_dict,varscan=True):
    '''
    vtools select TCGA_CH_5748 'vsvcf.STATUS_INFO="Somatic"' -t somatic
    vtools use ../vsvcf.ann --files ../TCGA-EJ-5514.somatic.snps.vars.raw.2.vcf ../TCGA-CH-5748.somatic.snps.vars.raw.2.vcf
   '''    

    for sp,file in sample_dict.iteritems():
        
        os.chdir(rootdir)
        args=['mkdir',sp]
        retcode = subprocess.call(args)
        os.chdir(os.path.join(rootdir,sp))
            
        args=['vtools','init',sp]
        retcode=subprocess.call(args)
    
        if retcode!=0:
            return JOB_ERROR
        cmd=['vtools','import']
        args=cmd+[file]+['--sample_name',sp]+['--build',build]
        
        if retcode != 0:
            raise JOB_ERROR
        else:
            print args
            retcode=subprocess.call(args)
    
        # load the varscan-vcf annotation.
        if retcode!=0:
            return JOB_ERROR
        elif varscan:
            args=['vtools','use',vsannfile,'--files']+[file]
            retcode=subprocess.call(args)
            if retcode != 0:
                raise JOB_ERROR        
            retcode = load_all_annotation(rootdir,sp, annotation_dict)
        
       
        if retcode!=0:
            return JOB_ERROR
        else:
            f= os.path.join(rootdir,sp)
            outfile = os.path.join(f,'somatic.vars.txt')
            retcode = report_somatic_mutations(rootdir,sp,outfile)
            if retcode!=0:
                return JOB_ERROR

        
    return JOB_SUCCESS



def find_somatic_mutations(rootdir,name,outfile,table=None):
    '''
    create a table of all somatic mutations observed accross patients
    load some annotation databases, gene annotion and cancerGenes, annovar
    load NSdnSNP
    select out everything that is in dbSNP
    report. 
        vtools select variant 'vsvcf.STATUS_INFO="Somatic"' -t somatic3
    
    vtools output somatic3 chr pos ref alt ccdsGene.name vsvcf.STATUS_INFO vsvcf.SPV_INFO vsvcf.DPN1_INFO vsvcf.DPN2_INFO vsvcf.DPT1_INFO vsvcf.DPT2_INFO vsvcf.VFT_INFO -l 10
    $vtools output somatic3 chr pos ref alt ccdsGene.name refGene.name2 genename dbNSFP.aaref dbNSFP.aaalt SIFT_score PolyPhen2_score 
    PolyPhen2_score vsvcf.STATUS_INFO vsvcf.SPV_INFO vsvcf.DPN1_INFO vsvcf.DPN2_INFO vsvcf.DPT1_INFO vsvcf.DPT2_INFO vsvcf.VFT_INFO -l 10 --header
    '''
    # Change to folder with varianttools project
    filter=True
    annovarpath='/data/projects/tcgatest/vtools_annotation_files/annovar/'
    os.chdir( os.path.join(rootdir,name) )
    
    if table is not None:
        table='variant'

    table_somatic='%s_%s'%(table,'somatic')
    new_somatic='%s_%s'%(table,'somaticnonDBSNP')
    final_somatic='%s_%s'%(table,'somatic_report')
    final_somatic_exons =  '%s_%s'%(final_somatic,'nonsyn_exons')
    VTFrac,dp_vnormal, db_vtumor = 0.1,0,6
    if filter:
        thresholding = '\'vsvcf.STATUS_INFO=\"Somatic\" AND vsvcf.VFT_INFO > %f AND vsvcf.DPN2_INFO=%d AND vsvcf.DPT2_INFO > %d \''
        thresholding = [thresholding%(VTFrac,dp_vnormal, db_vtumor)]
    else:
        thresholding=['\'vsvcf.STATUS_INFO=\"Somatic\"\'']
    
    args=['vtools','select', 'variant']+thresholding+['-t', table_somatic]
    args=",".join(args).replace(',',' ')
    print args
    retcode=subprocess.call(args,shell=True)
    if retcode!=0:
        return JOB_ERROR
    args=['vtools','exclude',table_somatic,'\"dbSNP.chr is not NULL or thousandGenomes.chr is not NULL\"','-t',final_somatic]
    args=",".join(args).replace(',',' ')
    print args
    retcode=subprocess.call(args,shell=True)
    if retcode!=0:
        return JOB_ERROR
    args= ['vtools', 'export', final_somatic, 'ann_input', '--format', 'ANNOVAR']
    print args
    retcode=subprocess.call(args)
    if retcode!=0:
        return JOB_ERROR
    annovar=os.path.join(annovarpath,'annotate_variation.pl')
    humandb=os.path.join(annovarpath,'humandb/')
    hgbuild='hg19'
    args = [annovar, 'ann_input',humandb, '--build',hgbuild]
    print args
    retcode=subprocess.call(args)
    if retcode!=0:
        return JOB_ERROR

    args = ['vtools', 'update',final_somatic, '--from_file','ann_input.exonic_variant_function',
             '--format', 'ANNOVAR_output', '--var_info', 'mut_type', 'genename', 'function']
    print args
    retcode=subprocess.call(args)
    if retcode!=0:
        return JOB_ERROR
    
    args=['vtools','select','variant','\'function is not NULL AND mut_type is not "synonymous SNV"\'', '-t',final_somatic_exons]    
    args=",".join(args).replace(',',' ')
    print args
    retcode=subprocess.call(args,shell=True)
    if retcode!=0:
        return JOB_ERROR

    
    fields=['refGene.name2','chr', 'pos','ref', 'alt',
            'variant.mut_type', 'variant.function',
            'vsvcf.STATUS_INFO','vsvcf.SPV_INFO', 'vsvcf.DPN1_INFO', 'vsvcf.DPN2_INFO',
             'vsvcf.DPT1_INFO', 'vsvcf.DPT2_INFO', 'vsvcf.VFT_INFO',
             'CancerGeneCensus.CancerSomaticMut','TumourTypesSomatic',
             'KgDesc',
             'SIFT_score','Polyphen2_score','MutationTaster_pred',
             '-g', 'variant.variant_id']
    header=['--header','GENNAME','CHR','POS','REF','ALT','MUTATION_TYPE','LOCATION',
            'STATUS','PVAL','NORMAL_READS_REF','NORMAL_READS_ALT',
            'TUMOR_READS_REF','TUMOR_READS_ALT','TUMOR_VARIANT_FRACTION',
            'CANCER_SOMATIC_MUT','TUMOR_TYPE_SOMATIC','KEGG_PATH',
            'SIFT_SCORE','POLYPHEN2_SCORE', 'MUTATION_TASTER_PRED']
    args=['vtools','output', final_somatic_exons]+fields+ header +['|','uniq']
    args=",".join(args).replace(',',' ')
    print args
    f=open(outfile,'w')
    retcode=subprocess.call(args,stdout=f,shell=True)
    f.close()
    if retcode!=0:
        return JOB_ERROR
    
    return JOB_SUCCESS


def create_sample_tables(rootdir,name,samples_dict):
    '''
    vtools select variant --samples 'sample_name="TCGA-EJ-5514"' -t TCGA_EJ_5514
    ''' 
    os.chdir( os.path.join(rootdir,name) )
        
    for sp in samples_dict:
        args=['vtools', 'select', 'variant', '--samples', '\'sample_name=\"TCGA-EJ-5514\"', '-t', sp]
        retcode=subprocess.call(args)
        if retcode!=0:
            print 'sample %s failed when generating the sample table'%(sp)
            return JOB_ERROR

    return JOB_SUCCESS
    
def report_somatic_mutations(rootdir,name,outfile,sample_dict=[]):
    '''
    '''
    os.chdir( os.path.join(rootdir,name) )
    
    if not sample_dict:
        outfile='%s_%s'%(os.path.join(rootdir,outfile),'all')
        retcode = find_somatic_mutations(rootdir,name,outfile,name)
        if retcode!=0:
            return JOB_ERROR
        
    else:
        for sp,spfile in sample_dict.iteritems():
            table=sp
            outfile='%s_%s'%(spfile,'somatic_report.txt')
            retcode = find_somatic_mutations(rootdir,name,outfile,table)
            if retcode!=0:
                return JOB_ERROR

    return JOB_SUCCESS


def load_gene_annotations(rootdir,name,annotation_dict):
    '''
    This will load the gene annotation
    '''
    os.chdir(os.path.join(rootdir,name))
    annolist=['refGene','ccdsGene','ccdsGene_exon']
    cmd=['vtools','use']
    retcode=0
    for ga in annolist:
        if ga not in annotation_dict:
            continue
        args=cmd+[annotation_dict[ga]]
        if retcode != 0:
            raise JOB_ERROR
        else:
            print args
            retcode=subprocess.call(args)

    if retcode!=0:
        return JOB_ERROR
    return JOB_SUCCESS

def load_variant_annoations(rootdir,name,annotation_dict):
    '''
    'evs'=exome variant server (2500 exomes from the NHLBI),
    'evs_5400=exome variant server (5400 exomes from the NHLBI)'
    'dbSNP',
    'dbNSFP'=non-synonymous variants of CCDS genes,
    'thousandGenomes=1000 genomes'
    '''
    
    os.chdir(os.path.join(rootdir,name))
    annolist=['dbSNP','thousandGenomes'] #'evs','evs_5400',
    annolist2=['dbNSFP']
    cmd=['vtools','use']
    retcode=0
    for ga in annolist:
        args=cmd+[annotation_dict[ga]]
        if retcode != 0:
            raise JOB_ERROR
        else:
            print ",".join(args).replace(',', ' ')
            retcode=subprocess.call(args)

    for ga in annolist2:
        args=cmd+[annotation_dict[ga]]
        if retcode != 0:
            raise JOB_ERROR
        else:
            print ",".join(args).replace(',', ' ')
            retcode=subprocess.call(args)
    
    if retcode!=0:
        return JOB_ERROR
    return JOB_SUCCESS
    

def load_functional_annoations(rootdir,name,annotation_dict):
    '''
    two main annotation= kegg pathways
    and cancer genome census
    they both required to have loaded refGene or other gene annotation first
    '''
    
    os.chdir(os.path.join(rootdir,name))
    annolist={'CancerGeneCensus':'refGene.name2','keggPathway':'ccdsGene.name'}#'ccdsGene.name'} # test if you can link kegg pathways using the refGene.name2
    retcode=0
    for ga,linkby in annolist.iteritems():
        args=['vtools','use',annotation_dict[ga],'--linked_by',linkby] # refGene.name2 is the hugo symbol
        if retcode != 0:
            raise JOB_ERROR
        else:
            print args
            retcode=subprocess.call(args)

    if retcode!=0:
        return JOB_ERROR
    return JOB_SUCCESS


def load_all_annotation(rootdir,name,annotation_dict):
    '''
    This step is slow, but it only needs to be run once by project. 
    '''
    retcode = load_gene_annotations(rootdir,name,annotation_dict)
    if retcode==0:
        retcode = load_functional_annoations(rootdir,name,annotation_dict)
        if retcode==0:
            retcode = load_variant_annoations(rootdir,name,annotation_dict)
            if retcode!=0:
                return JOB_ERROR
            else:
                return JOB_SUCCESS

    
def read_files_folder(folderpath,ext):
    ''' '''
    # Read files in folder
    myfiles=defaultdict()
    for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
        filename=infile.split('/')[-1]
        sp=filename.split('.')[0]
        myfiles[sp]=os.path.join(folderpath,filename)
    return myfiles

            
'''
samples_dict={'fmpn63':'/exds/users/oabalbin/projects/fmpn/data/gatk_calls/FMPN63.gatk_snps.b.filtered.model.vcf',
'fmpn51':'/exds/users/oabalbin/projects/fmpn/data/gatk_calls/FMPN51.gatk_snps.b.filtered.model.vcf',
'fmpn54','/exds/users/oabalbin/projects/fmpn/data/gatk_calls/FMPN54.gatk_snps.b.filtered.model.vcf',
'fmpn52','/exds/users/oabalbin/projects/fmpn/data/gatk_calls/FMPN52.gatk_snps.b.filtered.model.vcf',
'fmpn53','/exds/users/oabalbin/projects/fmpn/data/gatk_calls/FMPN53.gatk_snps.b.filtered.model.vcf',
'fmpn65','/exds/users/oabalbin/projects/fmpn/data/gatk_calls/FMPN65.gatk_snps.b.filtered.model.vcf',
'fmpn62','/exds/users/oabalbin/projects/fmpn/data/gatk_calls/FMPN62.gatk_snps.b.filtered.model.vcf'}
'''


build='hg19'
rootdir='/data/projects/tcgatest/'
vcfdir='/data/projects/tcgatest/vcf_files2/'
vsannfile='/data/projects/tcgatest/vsvcf.ann'
varscan=True
#name='familial_fmpn'
name='tcga_PRAD2'
ext='snps.vars.raw.2.vcf'
annotation_dir = '/data/projects/tcgatest/vtools_annotation_files'
read_files_folder(rootdir,ext)

vtools_annotation_db = {'refGene':'refGene-hg19_20110909.DB','ccdsGene':'ccdsGene.DB',
                        'evs':'evs.DB','evs_5400':'evs_5400.DB','dbNSFP':'dbNSFP-hg18_hg19_1.1_2.DB',
                        'dbSNP':'dbSNP-hg19_132.DB','thousandGenomes':'thousandGenomes-hg19_20110909.DB',
                        'CancerGeneCensus':'CancerGeneCensus-20111215.DB','keggPathway':'keggPathway-20110823.DB',
                        } #'ccdsGene':'ccdsGene.DB','ccdsGene_exon':'ccdsGene_exon.DB',


for g, f in vtools_annotation_db.iteritems():
    vtools_annotation_db[g]=os.path.join(annotation_dir,f)

print vtools_annotation_db

#samples_dict = {'TCGA-EJ-5530':'/data/projects/tcgatest/TCGA-EJ-5530.somatic.snps.vars.raw.2.vcf'}


if __name__ == '__main__':
    
    samples_dict=read_files_folder(vcfdir,ext)
    print samples_dict
    args=['vtools','use',vsannfile,'--files']+samples_dict.values() # all snv files to create a giant annotation file of all mutations.
    print args
    create_children_project(rootdir,name,samples_dict,
                            build,vsannfile,vtools_annotation_db,varscan)
    #create_sample_tables(rootdir,name,samples_dict)
